#include "sysc/kernel/sc_time.h"
#include <random>
#include <systemc>
#include <tlm>
#include <tlm_utils/simple_initiator_socket.h>
#include <tlm_utils/simple_target_socket.h>

using namespace sc_core;
using namespace sc_dt;
using namespace tlm_utils;
using namespace std;

#define ADDR_WIDTH 8
#define DATA_WIDTH 32

SC_MODULE(Master) {
    // Random number generator
    std::mt19937 gen;
    std::uniform_int_distribution<uint8_t> addr_dist;
    std::uniform_int_distribution<uint8_t> rwop_dist; // 0: read, 1: write
    // Initiator socket
    simple_initiator_socket<Master> socket;

    SC_CTOR(Master) : socket("initiator_socket"),
                      addr_dist(0, 255),
                      rwop_dist(0, 1) {
        SC_THREAD(process);
    }

    void process() {
        sc_time delay = sc_time(10, SC_NS);
        while (true) {
            uint8_t addr = addr_dist(gen);
            uint8_t rwop = rwop_dist(gen);
            uint32_t data = 0;
            
            // Set up the payload
            tlm::tlm_generic_payload* payload = new tlm::tlm_generic_payload();
            tlm::tlm_command cmd = (rwop == 0) ? tlm::TLM_READ_COMMAND : tlm::TLM_WRITE_COMMAND;
            if (cmd == tlm::TLM_WRITE_COMMAND) { data = 0xFF000000 | addr; }
            payload->set_command(cmd);
            payload->set_address(addr);
            payload->set_data_ptr(reinterpret_cast<unsigned char*>(&data));
            payload->set_data_length(sizeof(data));
            payload->set_byte_enable_ptr(0);
            payload->set_dmi_allowed(false);
            payload->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE);

            // Send the transaction
            socket->b_transport(*payload, delay);

            // Check response status
            if (payload->is_response_error()) {
                SC_REPORT_ERROR("TLM-2", "Response error from b_transport");
            }

            printf("%s %s addr=0x%02X, data=0x%08X at time %s delay %s\n", 
                this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                addr, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            
            // Wait before next transaction
            wait(delay);
        }
    }
};

SC_MODULE(Slave) {
    // Target socket
    simple_target_socket<Slave> socket;
    // Memory
    uint32_t mem[256];

    SC_CTOR(Slave) : socket("target_socket") {
        socket.register_b_transport(this, &Slave::b_transport);
        // Initialize memory
        for (int i = 0; i < 256; ++i) {
            mem[i] = 0xAA000000 | i;
        }
    }

    void b_transport(tlm::tlm_generic_payload& payload, sc_time& delay) {
        // Extract command, address, and data pointer
        tlm::tlm_command cmd = payload.get_command();
        uint64_t addr = payload.get_address();
        unsigned char* data_ptr = payload.get_data_ptr();
        unsigned int data_length = payload.get_data_length();

        // Check address and data length
        if (addr >= 256 || data_length != 4) {
            payload.set_response_status(tlm::TLM_ADDRESS_ERROR_RESPONSE);
            return;
        }
        
        // Operate on memory
        uint32_t* data = reinterpret_cast<uint32_t*>(data_ptr);
        if (cmd == tlm::TLM_READ_COMMAND) {
            *data = mem[addr];
        } else if (cmd == tlm::TLM_WRITE_COMMAND) {
            mem[addr] = *data;
        }

        // Set response status
        payload.set_response_status(tlm::TLM_OK_RESPONSE);
    }
};

int sc_main(int argc, char **argv)
{
    // Instantiate modules
    Master master("master");
    Slave slave("slave");

    // Bind sockets
    master.socket.bind(slave.socket);

    // Start simulation
    sc_start(10000, SC_NS);
    return 0;
}
